"""
猫狗识别项目 - 集成学习模型训练脚本
包含数据处理、模型构建、训练、评估和保存功能
"""

import os
import numpy as np
from PIL import Image
import joblib
import faiss
from sklearn.ensemble import VotingClassifier, RandomForestClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
import warnings
warnings.filterwarnings('ignore')

# ==================== 1. 数据处理部分 ====================
print("=" * 50)
print("开始数据处理...")
print("=" * 50)

# 定义数据路径
DATA_DIR = r"c:\Users\xingn\Desktop\作业\cat_dog_data\data\train"
IMAGE_SIZE = (224, 224)  # 统一图片尺寸

def load_and_preprocess_images(data_dir, image_size=(224, 224)):
    """
    从指定文件夹读取猫狗图片数据
    
    参数:
        data_dir: 数据文件夹路径
        image_size: 目标图片尺寸 (宽, 高)
    
    返回:
        images: 图片数据数组
        labels: 标签数组 (0=猫, 1=狗)
    """
    images = []
    labels = []
    
    # 获取所有图片文件
    image_files = [f for f in os.listdir(data_dir) if f.endswith('.jpg')]
    print(f"找到 {len(image_files)} 张图片")
    
    # 处理每张图片
    for idx, filename in enumerate(image_files):
        try:
            # 读取图片
            img_path = os.path.join(data_dir, filename)
            img = Image.open(img_path)
            
            # 转换为RGB模式(处理灰度图)
            if img.mode != 'RGB':
                img = img.convert('RGB')
            
            # 调整图片大小为224x224像素
            img = img.resize(image_size)
            
            # 转换为numpy数组并归一化到[0, 1]
            img_array = np.array(img) / 255.0
            
            # 展平为一维向量(224*224*3 = 150528维)
            img_flat = img_array.flatten()
            
            images.append(img_flat)
            
            # 根据文件名提取标签: cat=0, dog=1
            if filename.startswith('cat'):
                labels.append(0)
            elif filename.startswith('dog'):
                labels.append(1)
            
            # 显示进度
            if (idx + 1) % 1000 == 0:
                print(f"已处理 {idx + 1}/{len(image_files)} 张图片")
                
        except Exception as e:
            print(f"处理图片 {filename} 时出错: {e}")
            continue
    
    # 转换为numpy数组
    images = np.array(images, dtype=np.float32)
    labels = np.array(labels, dtype=np.int32)
    
    print(f"\n数据处理完成!")
    print(f"图片数据形状: {images.shape}")
    print(f"标签数据形状: {labels.shape}")
    print(f"猫的数量: {np.sum(labels == 0)}")
    print(f"狗的数量: {np.sum(labels == 1)}")
    
    return images, labels

# 加载并预处理图片数据
X, y = load_and_preprocess_images(DATA_DIR, IMAGE_SIZE)

# 划分训练集和测试集(8:2比例)
print("\n" + "=" * 50)
print("划分训练集和测试集...")
print("=" * 50)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"训练集大小: {X_train.shape[0]} 样本")
print(f"测试集大小: {X_test.shape[0]} 样本")
print(f"特征维度: {X_train.shape[1]} 维")

# 特征标准化
print("\n对特征进行标准化...")
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)
print("标准化完成!")

# PCA降维(解决内存问题)
print("\n" + "=" * 50)
print("使用PCA进行降维...")
print("=" * 50)
print(f"原始特征维度: {X_train_scaled.shape[1]} 维")

# 降维到500维(保留大部分信息,同时大幅减少内存消耗)
n_components = 500
pca = PCA(n_components=n_components, random_state=42)
X_train_pca = pca.fit_transform(X_train_scaled)
X_test_pca = pca.transform(X_test_scaled)

explained_variance = np.sum(pca.explained_variance_ratio_) * 100
print(f"降维后特征维度: {X_train_pca.shape[1]} 维")
print(f"保留的方差比例: {explained_variance:.2f}%")
print("PCA降维完成!")

# ==================== 2. 构建Faiss索引 ====================
print("\n" + "=" * 50)
print("构建Faiss索引...")
print("=" * 50)

# 创建Faiss flat索引(L2距离)
dimension = X_train_scaled.shape[1]  # 特征维度
faiss_index = faiss.IndexFlatL2(dimension)

# 将训练数据添加到索引中
# Faiss要求float32类型
X_train_faiss = X_train_scaled.astype(np.float32)
faiss_index.add(X_train_faiss)

print(f"Faiss索引构建完成!")
print(f"索引中的向量数量: {faiss_index.ntotal}")

# ==================== 3. 集成学习模型构建 ====================
print("\n" + "=" * 50)
print("构建集成学习模型...")
print("=" * 50)

# 基础模型1: 随机森林分类器
# n_estimators: 树的数量
# max_depth: 限制树的最大深度,防止过拟合
# min_samples_split: 分裂内部节点所需的最小样本数
rf_classifier = RandomForestClassifier(
    n_estimators=100,       # 100棵树
    max_depth=20,           # 最大深度为20
    min_samples_split=10,   # 分裂需要至少10个样本
    min_samples_leaf=5,     # 叶节点至少5个样本
    random_state=42,
    n_jobs=-1               # 使用所有CPU核心
)
print("✓ 随机森林分类器配置完成")
print(f"  - 树的数量: 100")
print(f"  - 最大深度: 20")
print(f"  - 最小分裂样本数: 10")
print(f"  - 最小叶节点样本数: 5")

# 基础模型2: 逻辑回归分类器
# C: 正则化强度的倒数
# max_iter: 最大迭代次数
# solver: 优化算法
lr_classifier = LogisticRegression(
    C=1.0,                  # 正则化参数
    max_iter=1000,          # 最大迭代1000次
    solver='lbfgs',         # 使用L-BFGS优化算法
    random_state=42,
    n_jobs=-1               # 使用所有CPU核心
)
print("✓ 逻辑回归分类器配置完成")
print(f"  - 正则化参数C: 1.0")
print(f"  - 最大迭代次数: 1000")
print(f"  - 优化算法: L-BFGS")

# 构建投票集成分类器
# voting='hard': 硬投票 - 基于多数投票
# voting='soft': 软投票 - 基于预测概率的平均
print("\n构建投票集成分类器...")

# 硬投票集成
ensemble_model_hard = VotingClassifier(
    estimators=[
        ('random_forest', rf_classifier),
        ('logistic_regression', lr_classifier)
    ],
    voting='hard',  # 硬投票策略
    n_jobs=-1       # 并行训练
)
print("✓ 硬投票集成模型构建完成")

# 软投票集成
ensemble_model_soft = VotingClassifier(
    estimators=[
        ('random_forest', rf_classifier),
        ('logistic_regression', lr_classifier)
    ],
    voting='soft',  # 软投票策略
    n_jobs=-1       # 并行训练
)
print("✓ 软投票集成模型构建完成")

# ==================== 4. 模型训练与评估 ====================
print("\n" + "=" * 50)
print("开始模型训练...")
print("=" * 50)

# 训练硬投票模型(使用PCA降维后的数据)
print("\n[1/2] 训练硬投票集成模型...")
ensemble_model_hard.fit(X_train_pca, y_train)
print("✓ 硬投票模型训练完成")

# 评估硬投票模型
y_train_pred_hard = ensemble_model_hard.predict(X_train_pca)
y_test_pred_hard = ensemble_model_hard.predict(X_test_pca)

train_acc_hard = accuracy_score(y_train, y_train_pred_hard)
test_acc_hard = accuracy_score(y_test, y_test_pred_hard)
test_f1_hard = f1_score(y_test, y_test_pred_hard)

print(f"  训练集准确率: {train_acc_hard:.4f}")
print(f"  测试集准确率: {test_acc_hard:.4f}")
print(f"  测试集F1分数: {test_f1_hard:.4f}")

# 训练软投票模型(使用PCA降维后的数据)
print("\n[2/2] 训练软投票集成模型...")
ensemble_model_soft.fit(X_train_pca, y_train)
print("✓ 软投票模型训练完成")

# 评估软投票模型
y_train_pred_soft = ensemble_model_soft.predict(X_train_pca)
y_test_pred_soft = ensemble_model_soft.predict(X_test_pca)

train_acc_soft = accuracy_score(y_train, y_train_pred_soft)
test_acc_soft = accuracy_score(y_test, y_test_pred_soft)
test_f1_soft = f1_score(y_test, y_test_pred_soft)

print(f"  训练集准确率: {train_acc_soft:.4f}")
print(f"  测试集准确率: {test_acc_soft:.4f}")
print(f"  测试集F1分数: {test_f1_soft:.4f}")

# 对比两种投票策略
print("\n" + "=" * 50)
print("投票策略对比:")
print("=" * 50)
print(f"硬投票 - 测试集准确率: {test_acc_hard:.4f}, F1分数: {test_f1_hard:.4f}")
print(f"软投票 - 测试集准确率: {test_acc_soft:.4f}, F1分数: {test_f1_soft:.4f}")

# 选择性能更好的模型
if test_acc_soft >= test_acc_hard:
    ensemble_model = ensemble_model_soft
    print("\n✓ 选择软投票模型作为最终模型")
else:
    ensemble_model = ensemble_model_hard
    print("\n✓ 选择硬投票模型作为最终模型")

# ==================== 5. 模型保存 ====================
print("\n" + "=" * 50)
print("保存模型和索引...")
print("=" * 50)

# 创建保存目录
os.makedirs('./saved_models', exist_ok=True)
os.makedirs('./saved_indexes', exist_ok=True)

# 保存集成学习模型
# 使用joblib保存sklearn模型,保持模型的完整性
model_path = './saved_models/ensemble_model.joblib'
joblib.dump(ensemble_model, model_path)
print(f"✓ 集成模型已保存到: {model_path}")

# 保存标准化器(预测时需要使用相同的标准化)
scaler_path = './saved_models/scaler.joblib'
joblib.dump(scaler, scaler_path)
print(f"✓ 标准化器已保存到: {scaler_path}")

# 保存PCA模型(预测时需要使用相同的降维)
pca_path = './saved_models/pca.joblib'
joblib.dump(pca, pca_path)
print(f"✓ PCA模型已保存到: {pca_path}")

# 保存Faiss索引
# 使用Faiss自带的write_index函数保存索引
index_path = './saved_indexes/faiss_index.faiss'
faiss.write_index(faiss_index, index_path)
print(f"✓ Faiss索引已保存到: {index_path}")

# 保存训练集标签(用于Faiss检索时返回标签)
labels_path = './saved_indexes/train_labels.npy'
np.save(labels_path, y_train)
print(f"✓ 训练集标签已保存到: {labels_path}")

print("\n" + "=" * 50)
print("所有任务完成!")
print("=" * 50)
print("\n模型训练总结:")
print(f"- 训练样本数: {len(X_train)}")
print(f"- 测试样本数: {len(X_test)}")
print(f"- 最终模型测试准确率: {max(test_acc_hard, test_acc_soft):.4f}")
print(f"- 最终模型测试F1分数: {max(test_f1_hard, test_f1_soft):.4f}")
print("\n可以运行 ensemble_webapp.py 启动网页应用进行预测!")
